import os
import sys
sys.path.append('..')
import torch
import torch.nn as nn
import segmentation_models_pytorch as smp
from utils.json import json_to_image
from utils.helper import get_subdirs, set_logger, GPUManager, del_useless_folders
import argparse
import pandas as pd
import glob
import time
from pathlib import Path
import numpy as np
import cv2
import PIL
from datasets import T2_Seg_Dataset
import platform
import yaml
from utils.trainer import Seg_Trainer
from utils.helper import remove_dataparallel
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import logging

logger = logging.getLogger('SegModel.Train')
def parse_args():
    """
    Set args parameters
    """
    parser = argparse.ArgumentParser(description='Train a segmentation model.')
    parser.add_argument("--height", type=int, default=1024, help="size of image height")
    parser.add_argument("--width", type=int, default=1024, help="size of image width")
    parser.add_argument("--model", type=str, default='Unet', choices=['Unet', 'DeepLabV3Plus' 'PSPNet'], help="Define model name")
    parser.add_argument("--backbone", type=str, default='se_resnext50_32x4d', choices=['resnet50', 'se_resnext50_32x4d', 'mobilenet_v2', 'timm-mobilenetv3_large_100'], help="Define model name")
    parser.add_argument("--weight", type=str, default=r'/data2/autorepair/ruanzhifeng/autorepair_t7_10/code/adc_segmentation/ckpts/AA_T7_1890/0613_2044_Unet_se_resnext50_32x4d/Unet_se_resnext50_32x4d_best_iou.pth', help="Define the location of pretrained weights.")
    parser.add_argument("--save", type=str, default='../onnx',help="Define the save_dir.")
    parser.add_argument("--res", type=str, default='T7_1890_seg_20250615.onnx',help="Define the save name.")

    args = parser.parse_args()

    return args


def main():
    args = parse_args()

    os.makedirs(args.save, exist_ok=True)

    model = getattr(smp, args.model)(args.backbone, encoder_weights=None, classes=1, activation=None)

    print('weight: ', args.weight)
    print('res: ', args.res)
    params = torch.load(args.weight, map_location='cpu')
    try:
        params = params["state_dict"]
    except:
        params = params
    logger.info('Pretrained from %s has been loaded.' %args.weight)
    model.load_state_dict(remove_dataparallel(params))
    model.eval()
    
    dummy_input = torch.rand(1, 3, args.height, args.width)
    torch.onnx.export(model, dummy_input, os.path.join(args.save, args.res), input_names = ['input1'], output_names = ['output1'], opset_version=11)
    
if __name__ == "__main__":
    main()
